/*
* ov_detector.cpp
* Created on: 20230606
* Author: sumang
* Description: armor detector of openvino
*/

#include <unistd.h>
#include "ov_detect.h"

using namespace cv;
using namespace std;

void OvO_Detector::OvO_Init()
{
    cout << "load network" << endl;
    network_ = ie_.ReadNetwork(param_.xml_file_path, param_.bin_file_path);
    input_name_ = network_.getInputsInfo().begin()->first;
    for (auto iter : network_.getOutputsInfo())
    {
        auto dims = iter.second->getDims();
        int stride_h = param_.h / dims[2];
        int stride_w = param_.w / dims[3];
        assert(stride_h == stride_w && "Invalid stride!");
        output_layers_.push_back((s_OutLayer){
            .idx = (int)output_names_.size(),
            .stride = stride_h,
            .num_anchor = (int)dims[1],
            .num_out = (int)dims[4]});
        output_names_.push_back(iter.first);
        /*
        cout << dims[4] << endl;
        cout << param_.classes << endl;
        cout << param_.colors << endl;
        cout <<  param_.sizes << endl;
        */
        assert(dims[4] == 5 + 10 + param_.classes + param_.colors + param_.sizes && "Output dimension wrong!");
        iter.second->setPrecision(InferenceEngine::Precision::FP32);
    }
    executable_network_ = ie_.LoadNetwork(network_, "GPU");
    infer_request_ = executable_network_.CreateInferRequest();
    cout << "network_init_done. " << endl;
    vector<float> tmp(param_.w*param_.h*3);
    blob = tmp;
}

void OvO_Detector::copyBlob(vector<float> &blob, InferenceEngine::Blob::Ptr &ieBlob) {
    InferenceEngine::MemoryBlob::Ptr mblob =
        InferenceEngine::as<InferenceEngine::MemoryBlob>(ieBlob);
    if (!mblob) {
        THROW_IE_EXCEPTION
            << "We expect blob to be inherited from MemoryBlob in matU8ToBlob, "
            << "but by fact we were not able to cast inputBlob to MemoryBlob";
    }
    // locked memory holder should be alive all time while access to its buffer
    // happens
    auto mblobHolder = mblob->wmap();

    float *ie_blob_data = mblobHolder.as<float *>();

    memcpy(ie_blob_data, blob.data(), sizeof(float) * blob.size());
}

void OvO_Detector::preprocess()
{
    resize(input_img_, input_img_, Size(640,640));
    
    int img_h = input_img_.rows;
    int img_w = input_img_.cols;
    float *blob_data = blob.data();
    //cout << blob.size() << endl;
    size_t i = 0;
    for (size_t row = 0; row < img_h; ++row)
    {
        uchar *uc_pixel = input_img_.data + row * input_img_.step;
        for (size_t col = 0; col < img_w; ++col)
        {
            // 三通道
            blob_data[i] = (float)uc_pixel[2] / 255.0;
            blob_data[i + img_h * img_w] = (float)uc_pixel[1] / 255.0;
            blob_data[i + 2 * img_h * img_w] = (float)uc_pixel[0] / 255.0;
            uc_pixel += 3;
            ++i;
        }
    }
}
void OvO_Detector::inference()
{
    InferenceEngine::Blob::Ptr ieBlob = infer_request_.GetBlob(input_name_);
    copyBlob(blob, ieBlob);
    auto infer_start = std::chrono::steady_clock::now();
    infer_request_.Infer();
}
void OvO_Detector::postprocess()
{
    auto decode_start = std::chrono::steady_clock::now();
    for (auto layer : output_layers_)
    {
        const InferenceEngine::Blob::Ptr output_blob =
            infer_request_.GetBlob(output_names_[layer.idx]);
        InferenceEngine::MemoryBlob::CPtr moutput =
            InferenceEngine::as<InferenceEngine::MemoryBlob>(output_blob);
        auto moutputHolder = moutput->rmap();
        const float *net_pred =
            moutputHolder.as<const InferenceEngine::PrecisionTrait<
                InferenceEngine::Precision::FP32>::value_type *>();

        decode_outputs(net_pred, outputs_, layer, param_.w, param_.h);
    }
    do_merge_nms();
}
bool OvO_Detector::detect()
{
    if (!input_imgs_.empty())
    {
        // cout << "now detect" << endl;
        img_mutex_.lock();
        input_img_ = input_imgs_.back();
        input_imgs_.clear();
        img_mutex_.unlock();

        //auto infer_start = std::chrono::steady_clock::now();
        preprocess();
        //auto infer_end = std::chrono::steady_clock::now();
        //auto duration = std::chrono::duration<double, milli>(infer_end-infer_start).count();
        //cout << "preprocess: " << duration << endl;

        //auto infer_start = std::chrono::steady_clock::now();
        inference();
        //auto infer_end = std::chrono::steady_clock::now();
        //auto duration = std::chrono::duration<double, milli>(infer_end-infer_start).count();
        //cout << "inference: " << duration << endl;

        //infer_start = std::chrono::steady_clock::now();
        outputs_.armor.clear();
        postprocess();
        outputs_.armor = output_nms_;

        //infer_end = std::chrono::steady_clock::now();
        //duration = std::chrono::duration<double, milli>(infer_end-infer_start).count();
        //cout << "postprocess: " << duration << endl;
        /*
        auto start = std::chrono::steady_clock::now();
        for (int i=0;i<1000;i++)
        {
            auto infer_start = std::chrono::steady_clock::now();
            preprocess();
            auto infer_end = std::chrono::steady_clock::now();
            auto duration = std::chrono::duration<double, milli>(infer_end-infer_start).count();
            //cout << "preprocess: " << duration << endl;

            infer_start = std::chrono::steady_clock::now();
            inference();
            infer_end = std::chrono::steady_clock::now();
            duration = std::chrono::duration<double, milli>(infer_end-infer_start).count();
            //cout << "inference: " << duration << endl;

            infer_start = std::chrono::steady_clock::now();
            outputs_.armor.clear();
            postprocess();
            outputs_.armor = output_nms_;

            infer_end = std::chrono::steady_clock::now();
            duration = std::chrono::duration<double, milli>(infer_end-infer_start).count();
            //cout << "postprocess: " << duration << endl;


        }
        auto end = std::chrono::steady_clock::now();
        auto dur = std::chrono::duration<double, milli>(end-start).count();
        cout << "fps: " << 1000/(dur/1000) << endl;
        exit(1);
        */
    }
    else ;//cout << "empty" << endl;
    return 1;
}



void OvO_Detector::decode_outputs(const float *prob, s_detections &objects,
                              s_OutLayer& layer_info, const int img_w,
                              const int img_h)
{
    float scale = std::min(param_.w / (img_w * 1.0), param_.h / (img_h * 1.0));

    std::vector<float> *l_anchor = nullptr;
    switch (layer_info.stride)
    {
        case 4:
            l_anchor = &param_.a1;
            break;
        case 8:
            l_anchor = &param_.a2;
            break;
        case 16:
            l_anchor = &param_.a3;
            break;
        case 32:
            l_anchor = &param_.a4;
            break;
        default:
            assert(false && "Unknown layer stride!");
            break;
    }

    int out_h = param_.h / layer_info.stride;
    int out_w = param_.w / layer_info.stride;
    float pred_data[layer_info.num_out];
    int nums = 0;
    for (int na = 0; na < 3; ++na)
    {
        for (int h_id = 0; h_id < out_h; ++h_id)
        {
            for (int w_id = 0; w_id < out_w; ++w_id)
            {
                int data_idx = (na * out_h * out_w + h_id * out_w + w_id) *
                            layer_info.num_out;
                float obj_conf = sigmoid(prob[data_idx + 4]);
                if (obj_conf > param_.bbox_conf_thresh)
                {
                    // std::cout << obj_conf << std::endl;
                    sigmoid(prob + data_idx, pred_data, 5);
                    sigmoid(prob + data_idx + 15, pred_data + 15,
                            param_.classes + param_.colors + 2);
                    memcpy(pred_data + 5, prob + data_idx + 5,
                        sizeof(float) * 10);
                    int col_id = std::max_element(pred_data + 15 + param_.classes,
                                                pred_data + 15 + param_.classes +
                                                    param_.colors) -
                                (pred_data + 15 + param_.classes);
                    if (col_id == param_.camp) continue;
                    int cls_id =
                        std::max_element(pred_data + 15,
                                        pred_data + 15 + param_.classes) -
                        (pred_data + 15);
                    
                    int t_size = std::max_element(pred_data + 15 + param_.classes + param_.colors,
                            pred_data + 15 + param_.classes + param_.colors + 2) -
                    (pred_data + 15 + param_.classes + param_.colors);

                    double final_conf =
                        obj_conf * pow(pred_data[15 + cls_id] *
                                        pred_data[15 + param_.classes + col_id] *
                                        pred_data[15 + param_.classes + param_.colors + t_size],
                                        1/3.);
                    if (final_conf > param_.bbox_conf_thresh)
                    {
                        nums++;
                        // std::cout << final_conf << " " << col_id << " "
                        //           << cls_id << std::endl;
                        
                        s_armor now;
                        float x = (pred_data[0] * 2.0 - 0.5 + w_id) *
                                layer_info.stride;
                        float y = (pred_data[1] * 2.0 - 0.5 + h_id) *
                                layer_info.stride;
                        float w =
                            pow(pred_data[2] * 2, 2) * l_anchor->at(na * 2);
                        float h =
                            pow(pred_data[3] * 2, 2) * l_anchor->at(na * 2 + 1);

                        for (int p = 0; p < 5; ++p)
                        {
                            float px =
                                (pred_data[5 + p * 2] * l_anchor->at(na * 2) +
                                w_id * layer_info.stride) /
                                scale;
                            float py = (pred_data[5 + p * 2 + 1] *
                                            l_anchor->at(na * 2 + 1) +
                                        h_id * layer_info.stride) /
                                    scale;
                            px = std::max(std::min(px, (float)(img_w)), 0.f);
                            py = std::max(std::min(py, (float)(img_h)), 0.f);
                            now.pts[p] = Point2f(px, py);
                            // std::cout << px << " " << py  << " ";
                        }
                        // std::cout << std::endl;

                        float x0 = (x - w * 0.5) / scale;
                        float y0 = (y - h * 0.5) / scale;
                        float x1 = (x + w * 0.5) / scale;
                        float y1 = (y + h * 0.5) / scale;

                        x0 = std::max(std::min(x0, (float)(img_w)), 0.f);
                        y0 = std::max(std::min(y0, (float)(img_h)), 0.f);
                        x1 = std::max(std::min(x1, (float)(img_w)), 0.f);
                        y1 = std::max(std::min(y1, (float)(img_h)), 0.f);

                        now.x_c = (x0+x1)/2;
                        now.y_c = (y0+y1)/2;
                        now.rect = Rect(x0, y0, x1 - x0, y1 - y0);
                        now.conf = final_conf;
                        now.color = col_id;
                        now.type = cls_id;
                        now.t_size = t_size;
                        objects.armor.push_back(now);
                        
                    }
                }
            }
        }
    }
}


static float calc_iou(const s_armor& a,const s_armor& b){
    Rect_<float> inter = a.rect & b.rect;
    float inter_area = inter.area();
    float union_area = a.rect.area() + b.rect.area() - inter_area;
    return inter_area / union_area;
}


void OvO_Detector::do_merge_nms(){
    output_nms_.clear();
    std::vector<pick_merge_store> picked;

    std::sort(outputs_.armor.begin(),outputs_.armor.end(),armor_compare());
    for(int i = 0;i < outputs_.armor.size();++i){
        s_armor& now = outputs_.armor[i];
        bool keep = 1;
        for(int j = 0;j < picked.size();++j){
            s_armor& pre = outputs_.armor[picked[j].id];
            float iou = calc_iou(now,pre);
            
            //store for merge_nms
            if(iou > param_.nms_thresh || isnan(iou)){
                keep = 0;
                if(iou > param_.merge_thresh && now.color == pre.color && now.type == pre.type && now.t_size == pre.t_size){
                    picked[j].merge_confs.push_back(now.conf);
                    for(int k = 0;k < 5; ++k){
                        picked[j].merge_pts.push_back(now.pts[k]);
                    }
                }
                break;
            }
        }
        if(keep){
            picked.push_back({i,{},{}});
        }
    }
    for(int i = 0;i < picked.size();++i){
        int merge_num = picked[i].merge_confs.size();
        s_armor now = outputs_.armor[picked[i].id];
        double conf_sum = now.conf;
        for(int j = 0;j < 5;++j) now.pts[j] *= now.conf;
        for(int j = 0;j < merge_num;++j){
            for(int k = 0;k < 5;++k){
                now.pts[k] += picked[i].merge_pts[j * 5 + k] * picked[i].merge_confs[j];
            }
            conf_sum += picked[i].merge_confs[j];
        }
        for(int j = 0;j < 5;++j) now.pts[j] /= conf_sum;
        output_nms_.emplace_back(now);
    }
}


